// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/expression/cast_hooks.h>
#include <pollux/expression/function_call_to_special_form.h>
#include <pollux/expression/special_form.h>

namespace kumo::pollux::exec {
    constexpr melon::StringPiece kCast = "cast";
    constexpr melon::StringPiece kTryCast = "try_cast";

    /// Custom operator for casts from and to custom types.
    class CastOperator {
    public:
        virtual ~CastOperator() = default;

        /// Determines whether the cast operator supports casting to the custom type
        /// from the other type.
        virtual bool isSupportedFromType(const TypePtr &other) const = 0;

        /// Determines whether the cast operator supports casting from the custom type
        /// to the other type.
        virtual bool isSupportedToType(const TypePtr &other) const = 0;

        /// Casts an input vector to the custom type. This function should not throw
        /// when processing input rows, but report errors via context.setError().
        /// @param input The flat or constant input vector
        /// @param context The context
        /// @param rows Non-null rows of input
        /// @param resultType The result type.
        /// @param result The result vector of the custom type
        virtual void castTo(
            const BaseVector &input,
            exec::EvalCtx &context,
            const SelectivityVector &rows,
            const TypePtr &resultType,
            VectorPtr &result) const = 0;

        virtual void castTo(
            const BaseVector &input,
            exec::EvalCtx &context,
            const SelectivityVector &rows,
            const TypePtr &resultType,
            VectorPtr &result,
            const std::shared_ptr<CastHooks> & /* hooks */) const {
            castTo(input, context, rows, resultType, result);
        }

        /// Casts a vector of the custom type to another type. This function should
        /// not throw when processing input rows, but report errors via
        /// context.setError().
        /// @param input The flat or constant input vector
        /// @param context The context
        /// @param rows Non-null rows of input
        /// @param resultType The result type
        /// @param result The result vector of the destination type
        virtual void castFrom(
            const BaseVector &input,
            exec::EvalCtx &context,
            const SelectivityVector &rows,
            const TypePtr &resultType,
            VectorPtr &result) const = 0;
    };

    class CastExpr : public SpecialForm {
    public:
        /// @param type The target type of the cast expression
        /// @param expr The expression to cast
        /// @param trackCpuUsage Whether to track CPU usage
        CastExpr(
            TypePtr type,
            ExprPtr &&expr,
            bool trackCpuUsage,
            bool nullOnFailure,
            std::shared_ptr<CastHooks> hooks)
            : SpecialForm(
                  type,
                  std::vector<ExprPtr>({expr}),
                  nullOnFailure ? kTryCast.data() : kCast.data(),
                  false /* supportsFlatNoNullsFastPath */,
                  trackCpuUsage),
              nullOnFailure_(nullOnFailure),
              hooks_(std::move(hooks)) {
        }

        void evalSpecialForm(
            const SelectivityVector &rows,
            EvalCtx &context,
            VectorPtr &result) override;

        std::string toString(bool recursive = true) const override;

        std::string toSql(std::vector<VectorPtr> *) const override;

    private:
        /// Apply the cast after generating the input vectors
        /// @param rows The list of rows being processed
        /// @param input The input vector to be casted
        /// @param context The context
        /// @param fromType the input type
        /// @param toType the target type
        /// @param result The result vector
        void apply(
            const SelectivityVector &rows,
            const VectorPtr &input,
            exec::EvalCtx &context,
            const TypePtr &fromType,
            const TypePtr &toType,
            VectorPtr &result);

        VectorPtr applyMap(
            const SelectivityVector &rows,
            const MapVector *input,
            exec::EvalCtx &context,
            const MapType &fromType,
            const MapType &toType);

        VectorPtr applyArray(
            const SelectivityVector &rows,
            const ArrayVector *input,
            exec::EvalCtx &context,
            const ArrayType &fromType,
            const ArrayType &toType);

        VectorPtr applyRow(
            const SelectivityVector &rows,
            const RowVector *input,
            exec::EvalCtx &context,
            const RowType &fromType,
            const TypePtr &toType);

        /// Apply the cast between decimal vectors.
        /// @param rows Non-null rows of the input vector.
        /// @param input The input decimal vector. It is guaranteed to be flat or
        /// constant.
        template<typename ToDecimalType>
        VectorPtr applyDecimal(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType,
            const TypePtr &toType);

        // Apply the cast to a vector after vector encodings being peeled off. The
        // input vector is guaranteed to be flat or constant.
        void applyPeeled(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType,
            const TypePtr &toType,
            VectorPtr &result);

        template<typename Func>
        void applyToSelectedNoThrowLocal(
            EvalCtx &context,
            const SelectivityVector &rows,
            VectorPtr &result,
            Func &&func);

        /// The per-row level Kernel
        /// @tparam ToKind The cast target type
        /// @tparam FromKind The expression type
        /// @tparam TPolicy The policy used by the cast
        /// @param row The index of the current row
        /// @param input The input vector (of type FromKind)
        /// @param result The output vector (of type ToKind)
        template<TypeKind ToKind, TypeKind FromKind, typename TPolicy>
        void applyCastKernel(
            vector_size_t row,
            EvalCtx &context,
            const SimpleVector<typename TypeTraits<FromKind>::NativeType> *input,
            FlatVector<typename TypeTraits<ToKind>::NativeType> *result);

        VectorPtr castFromDate(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &toType);

        VectorPtr castToDate(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType);

        VectorPtr castFromIntervalDayTime(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &toType);

        template<typename TInput, typename TOutput>
        void applyDecimalCastKernel(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType,
            const TypePtr &toType,
            VectorPtr &castResult);

        template<typename TInput, typename TOutput>
        void applyIntToDecimalCastKernel(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &toType,
            VectorPtr &castResult);

        template<typename TInput>
        VectorPtr applyIntToBinaryCast(
            const SelectivityVector &rows,
            exec::EvalCtx &context,
            const BaseVector &input);

        template<typename TInput, typename TOutput>
        void applyFloatingPointToDecimalCastKernel(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &toType,
            VectorPtr &castResult);

        template<typename T>
        void applyVarcharToDecimalCastKernel(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &toType,
            VectorPtr &castResult);

        template<typename FromNativeType, TypeKind ToKind>
        VectorPtr applyDecimalToFloatCast(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType,
            const TypePtr &toType);

        template<typename FromNativeType, TypeKind ToKind>
        VectorPtr applyDecimalToIntegralCast(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType,
            const TypePtr &toType);

        template<typename FromNativeType>
        VectorPtr applyDecimalToBooleanCast(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context);

        template<typename FromNativeType>
        VectorPtr applyDecimalToPrimitiveCast(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType,
            const TypePtr &toType);

        template<TypeKind ToKind, TypeKind FromKind>
        void applyCastPrimitives(
            const SelectivityVector &rows,
            exec::EvalCtx &context,
            const BaseVector &input,
            VectorPtr &result);

        template<typename FromNativeType>
        VectorPtr applyDecimalToVarcharCast(
            const SelectivityVector &rows,
            const BaseVector &input,
            exec::EvalCtx &context,
            const TypePtr &fromType);

        template<TypeKind ToKind>
        void applyCastPrimitivesDispatch(
            const TypePtr &fromType,
            const TypePtr &toType,
            const SelectivityVector &rows,
            exec::EvalCtx &context,
            const BaseVector &input,
            VectorPtr &result);

        VectorPtr applyTimestampToVarcharCast(
            const TypePtr &toType,
            const SelectivityVector &rows,
            exec::EvalCtx &context,
            const BaseVector &input);

        bool nullOnFailure() const {
            return nullOnFailure_;
        }

        bool setNullInResultAtError() const {
            return nullOnFailure() && inTopLevel;
        }

        CastOperatorPtr getCastOperator(const TypePtr &type);

        // Custom cast operators for to and from top-level as well as nested types.
        melon::F14FastMap<std::string, CastOperatorPtr> castOperators_;

        bool nullOnFailure_;
        std::shared_ptr<CastHooks> hooks_;

        bool inTopLevel = false;
    };

    class CastCallToSpecialForm : public FunctionCallToSpecialForm {
    public:
        TypePtr resolveType(const std::vector<TypePtr> &argTypes) override;

        ExprPtr constructSpecialForm(
            const TypePtr &type,
            std::vector<ExprPtr> &&compiledChildren,
            bool trackCpuUsage,
            const core::QueryConfig &config) override;
    };

    class TryCastCallToSpecialForm : public FunctionCallToSpecialForm {
    public:
        TypePtr resolveType(const std::vector<TypePtr> &argTypes) override;

        ExprPtr constructSpecialForm(
            const TypePtr &type,
            std::vector<ExprPtr> &&compiledChildren,
            bool trackCpuUsage,
            const core::QueryConfig &config) override;
    };
} // namespace kumo::pollux::exec

#include <pollux/expression/cast_expr-inl.h>
